import numpy as np
import torch

from .hand import HandModel, RIGHT_HAND_INDEX
from .hand_skinning import skin_landmarks, _skin_points, _get_skinning_weights
from .handpose import SingleHandPose


def skin_landmarks_np(
        hand_model: HandModel,
        joint_angles: np.ndarray,
        wrist_transforms: np.ndarray,
) -> np.ndarray:
    landmarks = skin_landmarks(
        hand_model,
        torch.from_numpy(joint_angles).float(),
        torch.from_numpy(wrist_transforms).float(),
    )
    return landmarks.numpy()


def landmarks_from_hand_pose(
        hand_model: HandModel, hand_pose: SingleHandPose, hand_idx: int
) -> np.ndarray:
    """
    Compute 3D landmarks in the world space given the hand model and hand pose.
    """
    xf = hand_pose.wrist_xform.copy()
    # This function expects the user hand model to be a left hand.
    if hand_idx == RIGHT_HAND_INDEX:
        xf[:, 0] *= -1
    landmarks = skin_landmarks_np(hand_model, hand_pose.joint_angles, xf)
    return landmarks


def all_vertices_from_hand_pose(
        hand_model: HandModel, hand_pose: SingleHandPose, hand_idx: int
) -> np.ndarray:
    """
    计算手部模型所有网格顶点在世界坐标系中的位置

    Args:
        hand_model: 手部模型
        hand_pose: 手部姿态
        hand_idx: 手部索引（0表示右手，1表示左手）

    Returns:
        vertices: 所有网格顶点在世界坐标系中的位置 [V, 3]
    """
    if hand_model.mesh_vertices is None or hand_model.dense_bone_weights is None:
        raise ValueError("手部模型缺少网格数据或蒙皮权重")

    # 确保所有输入张量都有 batch 维度
    if len(hand_model.joint_rotation_axes.shape) == 2:  # 如果还没有batch维度
        # keep a reference to the original hand_model data while we create a new batched HandModel
        _orig_hand_model = hand_model
        # prepare hand_scale safely: if it's a tensor, clone/detach; if it's a python number, convert
        if _orig_hand_model.hand_scale is not None:
            if isinstance(_orig_hand_model.hand_scale, torch.Tensor):
                hand_scale_param = _orig_hand_model.hand_scale.clone().detach().unsqueeze(0)
            else:
                hand_scale_param = torch.tensor(_orig_hand_model.hand_scale).unsqueeze(0)
        else:
            hand_scale_param = None

        hand_model = HandModel(
            joint_rotation_axes=hand_model.joint_rotation_axes.unsqueeze(0),  # [22, 3] -> [1, 22, 3]
            joint_rest_positions=hand_model.joint_rest_positions.unsqueeze(0),  # [22, 3] -> [1, 22, 3]
            joint_frame_index=hand_model.joint_frame_index.unsqueeze(0),
            joint_parent=hand_model.joint_parent.unsqueeze(0),
            joint_first_child=hand_model.joint_first_child.unsqueeze(0),
            joint_next_sibling=hand_model.joint_next_sibling.unsqueeze(0),
            landmark_rest_positions=hand_model.landmark_rest_positions.unsqueeze(0),  # [21, 3] -> [1, 21, 3]
            landmark_rest_bone_weights=hand_model.landmark_rest_bone_weights.unsqueeze(0),  # [21, 3] -> [1, 21, 3]
            landmark_rest_bone_indices=hand_model.landmark_rest_bone_indices.unsqueeze(0),  # [21, 3] -> [1, 21, 3]
            hand_scale=hand_scale_param,
            mesh_vertices=hand_model.mesh_vertices.unsqueeze(0) if hand_model.mesh_vertices is not None else None,
            # [788, 3] -> [1, 788, 3]
            mesh_triangles=hand_model.mesh_triangles.unsqueeze(0),
            dense_bone_weights=hand_model.dense_bone_weights.unsqueeze(
                0) if hand_model.dense_bone_weights is not None else None,  # [788, 17] -> [1, 788, 17]
            joint_limits=hand_model.joint_limits.unsqueeze(0) if hand_model.joint_limits is not None else None
        )

    # 处理手腕变换矩阵
    xf = hand_pose.wrist_xform.copy()
    if hand_idx == RIGHT_HAND_INDEX:
        xf[:, 0] *= -1

    # 转换为张量并添加批次维度
    joint_angles = torch.from_numpy(hand_pose.joint_angles).float().unsqueeze(0)  # [20] -> [1, 20]
    wrist_transform = torch.from_numpy(xf).float().unsqueeze(0)  # [4, 4] -> [1, 4, 4]

    # 创建蒙皮权重矩阵
    skin_mat = _get_skinning_weights(
        torch.arange(17, device=hand_model.dense_bone_weights.device)
        .unsqueeze(0)
        .unsqueeze(0)
        .expand(1, hand_model.dense_bone_weights.shape[1], -1),  # [1, V, 17]
        hand_model.dense_bone_weights,  # [1, V, 17]
        17  # NUM_JOINT_FRAMES
    )

    # 计算变换后的顶点位置
    vertices = _skin_points(
        hand_model.joint_rest_positions,
        hand_model.joint_rotation_axes,
        skin_mat,
        joint_angles,
        hand_model.mesh_vertices,
        wrist_transform
    )

    # 移除批次维度并转换为 numpy 数组
    vertices = vertices.squeeze(0).detach().cpu().numpy()  # [1, V, 3] -> [V, 3]

    return vertices


def landmarks_from_batch_hand_pose(
        hand_model: HandModel,
        batch_joint_angles: torch.Tensor,  # shape (B, NUM_JOINTS_PER_HAND)
        batch_wrist_xforms: torch.Tensor,  # shape (B, 4, 4)
        batch_hand_idx: torch.Tensor  # shape (B,), values 0 or 1
) -> torch.Tensor:
    """
    批量计算一组手势的 3D landmarks，全部使用 PyTorch。
    Args:
        hand_model: HandModel，无 batch 维度
        batch_joint_angles: (B, NUM_JOINTS_PER_HAND)
        batch_wrist_xforms: (B, 4, 4)
        batch_hand_idx: (B,), RIGHT_HAND_INDEX 表示右手

    Returns:
        torch.Tensor: (B, NUM_LANDMARKS_PER_HAND, 3)
    """
    # 复制变换矩阵，避免原地修改
    xforms = batch_wrist_xforms.clone()

    # 对右手样本做 X 轴镜像
    right_mask = (batch_hand_idx == RIGHT_HAND_INDEX)  # shape (B,)
    xforms[right_mask, :, 0] *= -1  # 仅修改 x 轴

    # 直接调用 skin_landmarks，它支持 batch 的 tensor 输入
    landmarks = skin_landmarks(hand_model, batch_joint_angles, xforms)
    return landmarks  # shape: (B, NUM_LANDMARKS_PER_HAND, 3)


def all_vertices_from_batch_hand_pose(
        hand_model: HandModel,
        batch_joint_angles: torch.Tensor,  # shape (B, NUM_JOINTS_PER_HAND)
        batch_wrist_xforms: torch.Tensor,  # shape (B, 4, 4)
        batch_hand_idx: torch.Tensor  # shape (B,), values 0 or 1
) -> torch.Tensor:
    """
    批量计算一组手势的所有网格顶点，使用 PyTorch 实现。

    Args:
        hand_model: 手部模型（无 batch 维度）
        batch_joint_angles: 关节角度 (B, 20)
        batch_wrist_xforms: 手腕变换矩阵 (B, 4, 4)
        batch_hand_idx: 手部索引 (B,), 0=右手, 1=左手

    Returns:
        torch.Tensor: 所有网格顶点位置 (B, V, 3)
    """
    # 检查模型是否包含必要的网格数据
    if hand_model.mesh_vertices is None or hand_model.dense_bone_weights is None:
        raise ValueError("手部模型缺少网格数据或蒙皮权重")

    B = batch_joint_angles.size(0)  # 批量大小
    device = batch_joint_angles.device

    # 复制变换矩阵避免修改原始数据
    xforms = batch_wrist_xforms.clone()

    # 对右手样本进行X轴镜像处理
    right_mask = (batch_hand_idx == RIGHT_HAND_INDEX)  # 右手掩码 (B,)
    xforms[right_mask, :, 0] *= -1  # 仅修改X轴

    # 扩展模型参数以匹配批量大小
    def expand_tensor(t):
        if t is None:
            return None
        if t.dim() == 2:  # 无batch维度
            return t.unsqueeze(0).expand(B, *t.shape)
        return t

    # 准备模型参数
    joint_rest_positions = expand_tensor(hand_model.joint_rest_positions)  # (B, 22, 3)
    joint_rotation_axes = expand_tensor(hand_model.joint_rotation_axes)  # (B, 22, 3)
    mesh_vertices = expand_tensor(hand_model.mesh_vertices)  # (B, V, 3)
    dense_bone_weights = expand_tensor(hand_model.dense_bone_weights)  # (B, V, 17)

    # 创建关节索引张量
    joint_indices = torch.arange(17, device=device)  # (17,)
    joint_indices = joint_indices.unsqueeze(0).unsqueeze(0)  # (1, 1, 17)
    joint_indices = joint_indices.expand(B, dense_bone_weights.size(1), -1)  # (B, V, 17)

    # 计算蒙皮权重矩阵
    skin_mat = _get_skinning_weights(
        joint_indices,
        dense_bone_weights,
        17  # 关节数量
    )  # (B, V, 17)

    # 计算变换后的顶点位置
    vertices = _skin_points(
        joint_rest_positions,
        joint_rotation_axes,
        skin_mat,
        batch_joint_angles,
        mesh_vertices,
        xforms
    )  # (B, V, 3)

    return vertices

